# data/tushare_api.py
import pandas as pd
import tushare as ts

from config.config import config


class TushareClient:
    def __init__(self):
        # 初始化Tushare接口
        self.pro = ts.pro_api(config.get('tushare.token'))

    def get_daily_kline(self, ts_code, trade_date=None, start_date=None, end_date=None):
        """
        获取单只股票的日K线数据

        参数:
            trade_date (str): 获取某个交易日的数据（YYYYMMDD）
            stock_id (str): 股票代码（如：000001.SZ）
            start_date (str): 开始日期（YYYYMMDD）
            end_date (str): 结束日期（YYYYMMDD）

        返回:
            DataFrame: 包含日K数据的DataFrame
        """
        try:
            df = self.pro.daily(
                ts_code=ts_code,
                trade_date=trade_date,
                start_date=start_date,
                end_date=end_date
            )

            # 获取复权因子
            # adj_df = self.pro.adj_factor(ts_code=ts_code)
            # df = df.merge(adj_df, on='trade_date', how='left')

            # 获取股票基本信息
            # stock_info = self.pro.query('stock_basic', ts_code=stock_id)
            # if not stock_info.empty:
            #     df['symbol'] = stock_info['symbol'].values[0]
            #     df['name'] = stock_info['symbol'].values[0]

            return df

        except Exception as e:
            print(f"Tushare数据获取失败: {str(e)}")
            return pd.DataFrame()


if __name__ == '__main__':
    client = TushareClient()
    df = client.get_daily_kline('600610.SH', "20250704")
    print(df)
